package com.shayu.aicodemother.utils;

import cn.hutool.core.img.ImgUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import com.shayu.aicodemother.exception.BusinessException;
import com.shayu.aicodemother.exception.ErrorCode;
import io.github.bonigarcia.wdm.WebDriverManager;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.openqa.selenium.JavascriptExecutor;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeOptions;
import org.openqa.selenium.support.ui.WebDriverWait;

import java.io.File;
import java.time.Duration;
import java.util.UUID;

/**
 * 截图工具类（线程安全版本）
 */
@Slf4j
public class ThreadSaftyWebScreenshotUtils {

    // 使用 ThreadLocal 管理 WebDriver，每个线程有自己的驱动实例
    private static final ThreadLocal<WebDriver> driverThreadLocal = new ThreadLocal<>();
    
    // 配置参数
    private static final int DEFAULT_WIDTH = 1600;
    private static final int DEFAULT_HEIGHT = 900;
    private static final String IMAGE_SUFFIX = ".png";
    private static final String COMPRESS_SUFFIX = "_compressed.jpg";
    private static final float COMPRESSION_QUALITY = 0.3f;

    static {
        System.setProperty("wdm.timeout", "500");
        System.setProperty("wdm.retryCount", "3");
        // 设置国内镜像
        System.setProperty("wdm.chromeDownloadUrl", "https://npmmirror.com/mirrors/chromedriver/");
    }

    /**
     * 获取或创建 WebDriver 实例（线程安全）
     */
    private static WebDriver getDriver() {
        WebDriver driver = driverThreadLocal.get();
        if (driver == null) {
            driver = initChromeDriver(DEFAULT_WIDTH, DEFAULT_HEIGHT);
            driverThreadLocal.set(driver);
            log.debug("为线程 {} 创建新的 WebDriver 实例", Thread.currentThread().getName());
        }
        return driver;
    }

    /**
     * 清理当前线程的 WebDriver
     */
    public static void cleanupDriver() {
        WebDriver driver = driverThreadLocal.get();
        if (driver != null) {
            try {
                driver.quit();
                log.debug("线程 {} 的 WebDriver 已清理", Thread.currentThread().getName());
            } catch (Exception e) {
                log.warn("清理 WebDriver 时发生异常", e);
            } finally {
                driverThreadLocal.remove();
            }
        }
    }

    /**
     * 应用退出时清理所有 WebDriver
     */
    @PreDestroy
    public static void destroyAll() {
        // 清理所有线程的 WebDriver
        driverThreadLocal.remove();
        log.info("所有 WebDriver 实例已清理");
    }

    /**
     * 生成网页截图
     */
    public static String saveWebPageScreenshot(String webUrl) {
        if (StrUtil.isBlank(webUrl)) {
            log.error("网页截图失败，url为空");
            return null;
        }

        try {
            // 创建临时目录
            String rootPath = System.getProperty("user.dir") + "/tmp/screenshots/" + UUID.randomUUID().toString().substring(0, 8);
            FileUtil.mkdir(rootPath);
            
            // 生成文件路径
            String imageSavePath = rootPath + File.separator + RandomUtil.randomNumbers(5) + IMAGE_SUFFIX;
            String compressedImagePath = rootPath + File.separator + RandomUtil.randomNumbers(5) + COMPRESS_SUFFIX;

            // 获取当前线程的 WebDriver
            WebDriver driver = getDriver();
            
            // 访问网页并截图
            driver.get(webUrl);
            waitForPageLoad(driver);
            
            // 截图并保存
            byte[] screenshotBytes = ((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES);
            saveImage(screenshotBytes, imageSavePath);
            log.info("原始截图保存成功：{}", imageSavePath);
            
            // 压缩图片
            compressImage(imageSavePath, compressedImagePath);
            log.info("压缩图片保存成功：{}", compressedImagePath);
            
            // 删除原始图片
            FileUtil.del(imageSavePath);
            
            return compressedImagePath;
            
        } catch (Exception e) {
            log.error("网页截图失败：{}", webUrl, e);
            // 发生异常时清理当前线程的驱动
            cleanupDriver();
            return null;
        }
    }

    /**
     * 初始化 Chrome 浏览器驱动
     */
    private static WebDriver initChromeDriver(int width, int height) {
        try {
            WebDriverManager.chromedriver().useMirror().setup();
            
            ChromeOptions options = new ChromeOptions();
            options.addArguments(
                "--headless",
                "--disable-gpu",
                "--no-sandbox",
                "--disable-dev-shm-usage",
                "--disable-extensions",
                String.format("--window-size=%d,%d", width, height),
                "--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
            );

            WebDriver driver = new ChromeDriver(options);
            driver.manage().timeouts().pageLoadTimeout(Duration.ofSeconds(30));
            driver.manage().timeouts().implicitlyWait(Duration.ofSeconds(10));
            
            return driver;
        } catch (Exception e) {
            log.error("初始化 Chrome 浏览器失败", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "初始化 Chrome 浏览器失败");
        }
    }

    /**
     * 保存图片到文件
     */
    private static void saveImage(byte[] imageBytes, String imagePath) {
        try {
            FileUtil.writeBytes(imageBytes, imagePath);
        } catch (Exception e) {
            log.error("保存图片失败：{}", imagePath, e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "保存图片失败");
        }
    }

    /**
     * 压缩图片
     */
    private static void compressImage(String originImagePath, String compressedImagePath) {
        try {
            ImgUtil.compress(
                FileUtil.file(originImagePath),
                FileUtil.file(compressedImagePath),
                COMPRESSION_QUALITY
            );
        } catch (Exception e) {
            log.error("压缩图片失败：{} -> {}", originImagePath, compressedImagePath, e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "压缩图片失败");
        }
    }

    /**
     * 等待页面加载完成
     */
    private static void waitForPageLoad(WebDriver webDriver) {
        try {
            WebDriverWait wait = new WebDriverWait(webDriver, Duration.ofSeconds(10));
            wait.until(driver -> 
                ((JavascriptExecutor) driver)
                    .executeScript("return document.readyState")
                    .equals("complete")
            );
            Thread.sleep(2000);
            log.info("页面加载完成");
        } catch (Exception e) {
            log.error("等待页面加载时出现异常，继续执行截图", e);
        }
    }
}